import torch
import mindspore as ms

pt_path = "/root/linx/segan_pt/pt_disciminator.pth"
ms_path = "/root/linx/segan_ms/ms_disciminator.ckpt"

pt_ckpt = torch.load(pt_path, map_location='cpu')
ms_ckpt = ms.load_checkpoint(ms_path)
map_dict = {"act.weight": "act.w",
            "norm.bias": "norm.beta",
            "norm.weight": "norm.gamma",
            "norm.running_mean": "norm.moving_mean",
            "norm.running_var": "norm.moving_variance",
            "fc.1.weight": "fc.1.w", # prelu 参数
            "fc.3.weight": "fc.3.w", # prelu 参数
}
new_param_list = []
for name, param in pt_ckpt.items():
    new_name = 'D.' + name
    if "conv.weight" in new_name:
        param = param.unsqueeze(2)
    if "num_batches_tracked" in new_name:
        continue
    for old, new in map_dict.items():
        new_name = new_name.replace(old, new)
    new_param = ms.Tensor(param.numpy())
    new_param_list.append({"name": new_name, "data": new_param})

ms.save_checkpoint(new_param_list, "converted_d.ckpt")
pass
